import torch
from torch import nn
from basic_knowledge import d2lzh_pytorch as d2l

def corr2d_multi_in(X, K):
    # 沿着X和K的第0维（通道维）分别计算再相加
    res = d2l.corr2d(X[0, :, :], K[0, :, :])
    for i in range(1, X.shape[0]):
        res += d2l.corr2d(X[i, :, :], K[i, :, :])
    return res


X = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
              [[1, 2, 3], [4, 5, 6], [7, 8, 9]]])
K = torch.tensor([[[0, 1], [2, 3]], [[1, 2], [3, 4]]])

corr2d_multi_in(X, K)


# 实现一个互相关运算函数来计算多个通道的输出。
def corr2d_multi_in_out(X, K):
    # 对K的第0维遍历，每次同输入X做互相关计算。所有结果使用stack函数合并在一起
    return torch.stack([corr2d_multi_in(X, k) for k in K])


# 将核数组K同K+1（K中每个元素加一）和K+2连结在一起来构造一个输出通道数为3的卷积核。
K = torch.stack([K, K + 1, K + 2])
print(corr2d_multi_in_out(X, K))


# 使用全连接层中的矩阵乘法来实现1×1卷积。这里需要在矩阵乘法运算前后对数据形状做一些调整。
def corr2d_multi_in_out_1x1(X, K):
    c_i, h, w = X.shape  # c_i为输入的通道数
    c_o = K.shape[0]  # c_o为输入的通道数,即卷积核的数量
    X = X.view(c_i, h * w)  # 将三维向量改成二位向量，将每个通道展成一维，从这里可以看出来，pytorch里通道数在最前面
    K = K.view(c_o, c_i)  # K是4维的，c_o×3×1×1
    Y = torch.mm(K, X)  # 全连接层的矩阵乘法
    return Y.view(c_o, h, w)


X = torch.rand(3, 3, 3)
K = torch.rand(2, 3, 1, 1)

Y1 = corr2d_multi_in_out_1x1(X, K)
Y2 = corr2d_multi_in_out(X, K)
print(Y1.shape)
print(Y2.shape)  # 一样的

print((Y1 - Y2).norm().item() < 1e-6)
# 判断2个浮点数bai 是否相等，不du等用简单的用 a==b 判断，通常用 fabs(a-b)< eps 判断,
# 意思是当2数误差小于某 很小的数时，认为它们相等。例如 fabs(a-b)<1e-6.
# n=norm(A,p) 功能：norm函数可计算几种不同类型的矩阵范数,根据p的不同可得到不同的范数


# 1×1卷积层的作用与全连接层等价。
# 1×1卷积核的数量相当于输出的不同标签，通道数相当于参数数量，每个通道上的元素数量就是样本数量
# 相当于将同一个纵深的一个像素点（多个通道，相当于一个样本）经过一个卷积核的线性激活得到一个单一的结果
# 最终的结果是将每个样本的同一个线性模型的输出值放在一层通道上，每个通道就对应不同的卷积核（线性模型）